from random import random
from pylab import *
from numpy import arange

def logistic(a):
  return lambda x: a * x * (1 - x)

def iterate(f,start,n):
  ans = [start]
  for i in range(n):
    ans.append(f(ans[-1]))
  return ans

def attractors(f, start):
  orbit = iterate(f,start,500)
  return orbit[200:-1]

def plot_attractors(param,f,view,start):
  orbit = attractors(f, start)
  plot(list(param for x in orbit), list(view(x) for x in orbit), 'r,')

def bi_diagram(f_constructor, param_range,view,start):
  hold(True)
  for param in param_range:
    plot_attractors(param, f_constructor(param), view, start)
  show()

def logistic_diagram():
  bi_diagram(logistic,arange(1,4,0.01), lambda x: x, random())

def beetles(mu_a, state):
  """Lab Visit I from Alligood et. al., Chaos"""
  c_el = 0.012
  c_ea = 0.009
  c_pa = 0.004
  mu_l = 0.267
  mu_p = 0
  b = 7.48
  L, P, A = state
  L_new = b * A * exp((-c_ea * A) - (c_el * L))
  P_new = L * (1 - mu_l)
  A_new = P * (1 - mu_p) *  exp(-c_pa * A) + A * (1 - mu_a)
  return (L_new, P_new, A_new)

def beetle_diagram():
  print("Running beetle iterations...")
  start = (100, 100, 100)
  bi_diagram(lambda mu_a: (lambda s: beetles(mu_a, s)), arange(0,1,0.001), lambda s: s[0], start)

def main():
  def default():
    print('Usage: python '+sys.argv[0]+' [logistic|beetle]')
  try:
    {
      'logistic': logistic_diagram,
      'beetle': beetle_diagram,
    }.get(sys.argv[1], default)()
  except IndexError:
    default()

main()
